package hunyuan

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

type Client struct {
	SecretID  string
	SecretKey string
	Region    string
	Endpoint  string
	Service   string
	Version   string
	Algorithm string
}

func NewClient(secretID, secretKey, region string) *Client {
	return &Client{
		SecretID:  secretID,
		SecretKey: secretKey,
		Region:    region,
		Endpoint:  "ai3d.tencentcloudapi.com",
		Service:   "ai3d",
		Version:   "2025-05-13",
		Algorithm: "TC3-HMAC-SHA256",
	}
}

func (c *Client) tc3Sign(payload string, timestamp int64) (string, string, string) {
	date := time.Unix(timestamp, 0).UTC().Format("2006-01-02")
	
	// 1. Canonical Request
	payloadHash := sha256Hash(payload)
	canonicalRequest := fmt.Sprintf("POST\n/\n\ncontent-type:application/json\nhost:%s\n\ncontent-type;host\n%s",
		c.Endpoint, payloadHash)
	hashedCanonicalRequest := sha256Hash(canonicalRequest)

	// 2. String to Sign
	credentialScope := fmt.Sprintf("%s/%s/tc3_request", date, c.Service)
	stringToSign := fmt.Sprintf("%s\n%d\n%s\n%s",
		c.Algorithm, timestamp, credentialScope, hashedCanonicalRequest)

	// 3. Signature
	secretDate := hmacSHA256([]byte("TC3"+c.SecretKey), date)
	secretService := hmacSHA256(secretDate, c.Service)
	secretSigning := hmacSHA256(secretService, "tc3_request")
	signature := hex.EncodeToString(hmacSHA256(secretSigning, stringToSign))

	// 4. Authorization
	authorization := fmt.Sprintf("%s Credential=%s/%s, SignedHeaders=content-type;host, Signature=%s",
		c.Algorithm, c.SecretID, credentialScope, signature)

	return authorization, date, payloadHash
}

func (c *Client) request(action string, params map[string]interface{}) (map[string]interface{}, error) {
	url := "https://" + c.Endpoint + "/"
	timestamp := time.Now().Unix()
	
	payload, _ := json.Marshal(params)
	authorization, _, _ := c.tc3Sign(string(payload), timestamp)

	req, _ := http.NewRequest("POST", url, strings.NewReader(string(payload)))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Host", c.Endpoint)
	req.Header.Set("X-TC-Action", action)
	req.Header.Set("X-TC-Version", c.Version)
	req.Header.Set("X-TC-Timestamp", fmt.Sprintf("%d", timestamp))
	req.Header.Set("X-TC-Region", c.Region)
	req.Header.Set("Authorization", authorization)

	client := &http.Client{Timeout: 60 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	var result map[string]interface{}
	json.Unmarshal(body, &result)

	if response, ok := result["Response"].(map[string]interface{}); ok {
		if errObj, hasErr := response["Error"]; hasErr {
			errMap := errObj.(map[string]interface{})
			return nil, fmt.Errorf("[%s] %s", errMap["Code"], errMap["Message"])
		}
		return response, nil
	}

	return nil, errors.New("invalid response")
}

func (c *Client) SubmitJob(prompt, imageBase64, imageURL, format string, enablePBR bool) (string, error) {
	params := map[string]interface{}{
		"ResultFormat": format,
		"EnablePBR":    enablePBR,
	}

	count := 0
	if prompt != "" {
		params["Prompt"] = prompt
		count++
	}
	if imageBase64 != "" {
		params["ImageBase64"] = imageBase64
		count++
	}
	if imageURL != "" {
		params["ImageUrl"] = imageURL
		count++
	}

	if count != 1 {
		return "", errors.New("必须且只能提供 Prompt、ImageBase64、ImageUrl 之一")
	}

	resp, err := c.request("SubmitHunyuanTo3DRapidJob", params)
	if err != nil {
		return "", err
	}

	return resp["JobId"].(string), nil
}

type JobResult struct {
	Status       string
	ErrorCode    string
	ErrorMessage string
	ResultFiles  []ResultFile
}

type ResultFile struct {
	Type            string
	PreviewImageURL string
	URL             string
}

func (c *Client) QueryJob(jobID string) (*JobResult, error) {
	resp, err := c.request("QueryHunyuanTo3DRapidJob", map[string]interface{}{"JobId": jobID})
	if err != nil {
		return nil, err
	}

	result := &JobResult{
		Status:       resp["Status"].(string),
		ErrorCode:    getStringOrEmpty(resp, "ErrorCode"),
		ErrorMessage: getStringOrEmpty(resp, "ErrorMessage"),
	}

	if files, ok := resp["ResultFile3Ds"].([]interface{}); ok {
		for _, f := range files {
			fm := f.(map[string]interface{})
			result.ResultFiles = append(result.ResultFiles, ResultFile{
				Type:            fm["Type"].(string),
				PreviewImageURL: getStringOrEmpty(fm, "PreviewImageUrl"),
				URL:             fm["Url"].(string),
			})
		}
	}

	return result, nil
}

func (c *Client) WaitForDone(jobID string, interval, timeout int) ([]ResultFile, error) {
	start := time.Now()
	for {
		result, err := c.QueryJob(jobID)
		if err != nil {
			return nil, err
		}

		if result.Status == "DONE" {
			return result.ResultFiles, nil
		}
		if result.Status == "FAIL" {
			return nil, fmt.Errorf("任务失败: %s", result.ErrorMessage)
		}
		if time.Since(start).Seconds() > float64(timeout) {
			return nil, errors.New("任务超时")
		}

		time.Sleep(time.Duration(interval) * time.Second)
	}
}

func sha256Hash(s string) string {
	h := sha256.New()
	h.Write([]byte(s))
	return hex.EncodeToString(h.Sum(nil))
}

func hmacSHA256(key []byte, data string) []byte {
	h := hmac.New(sha256.New, key)
	h.Write([]byte(data))
	return h.Sum(nil)
}

func getStringOrEmpty(m map[string]interface{}, key string) string {
	if v, ok := m[key].(string); ok {
		return v
	}
	return ""
}

func DownloadFile(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return io.ReadAll(resp.Body)
}
